from abc import ABC, abstractmethod
from enum import Enum
from typing import Any, Protocol, TypeVar, Callable, Optional
import torch

# Type variables for better type hinting
T = TypeVar("T")  # Generic type for data


# Hook interfaces
class HookCallback(Protocol):
    """Protocol for hook callbacks.

    All hook callbacks should conform to this interface, which allows them
    to be called with variable arguments.
    """

    def __call__(self, *args: Any, **kwargs: Any) -> Any: ...


class BaseHookEvents(str, Enum):
    """Base class for hook events.

    This provides a foundation for defining standard hook points in the
    learning lifecycle.
    """

    pass


# Memory interfaces
class MemoryInterface(ABC):
    """Interface for memory management in incremental learning.

    Defines the contract for storing and retrieving examples for rehearsal.
    """

    @abstractmethod
    def add_examples(self, examples: Any, targets: Any) -> None:
        """Add examples to memory.

        Args:
            examples: Examples to store
            targets: Corresponding targets/labels
        """
        pass

    @abstractmethod
    def get_memory(self) -> Any:
        """Retrieve stored examples.

        Returns:
            Stored examples and their targets
        """
        pass

    @abstractmethod
    def update_memory(self, strategy: str = "random") -> None:
        """Update memory based on a specific strategy.

        Args:
            strategy: Strategy for updating memory (e.g., "random", "herding")
        """
        pass

    @abstractmethod
    def build_exemplars(
        self,
        data_manager: Any,
        network: Any,
        extract_features_func: Callable,
        known_classes: int,
        total_classes: int,
        device: Any,
    ) -> None:
        """Build exemplar set for rehearsal learning.

        Args:
            data_manager: Data manager for accessing dataset
            network: Neural network to extract features from
            extract_features_func: Function to extract features
            known_classes: Number of classes known from previous tasks
            total_classes: Total number of classes
            device: Device to perform computation on
        """
        pass


# Prototype interfaces
class PrototypeInterface(ABC):
    """Interface for prototype management.

    Defines the contract for computing and managing class prototypes.
    """

    @property
    @abstractmethod
    def feature_dim(self) -> int:
        """Return the computed class means."""
        pass

    @property
    @abstractmethod
    def num_classes(self) -> int:
        """Return the computed class means."""
        pass

    @property
    @abstractmethod
    def class_means(self) -> torch.Tensor:
        """Return the computed class means."""
        pass

    @property
    @abstractmethod
    def class_cov(self) -> torch.Tensor:
        """Return the computed class covariance."""
        pass

    @abstractmethod
    def fit(
        self,
        data_manager: Any,
        extract_tokens_fn: Callable,
        known_classes: int = 0,
        total_classes: int = 0,
        batch_size: int = 128,
        num_workers: int = 4,
        device: Optional[Any] = None,
    ) -> "PrototypeInterface":
        """Compute class statistics and create distributions.

        Args:
            data_manager: Data manager containing dataset
            extract_tokens_fn: Function to extract feature vectors
            known_classes: Number of previously known classes
            total_classes: Total number of classes including current task
            batch_size: Batch size for data loading
            num_workers: Number of worker processes for data loading
            device: Device to perform computation on

        Returns:
            Self for method chaining
        """
        pass
